#include "nn.h"

std::string name(grad& g) {
    std::string s;
    if(g.type == grad_type::DENSE) s += "dense";
    if(g.type == grad_type::LINCOMB) s += "lincomb (vars=" + std::to_string(((lincomb_grad&)g).vars) + ")";
    if(g.type == grad_type::LINCOMB_MODIFIED) s += "lincomb_modified (vars=" + std::to_string(((lincomb_modified_grad&)g).vars) + ")";
    if(g.type == grad_type::SCALAR) s += "scalar";
    s += " ins=" + std::to_string(g.in_size);
    s += " outs=" + std::to_string(g.out_size);
    return s;
}

grad* general_combine(grad& g, grad& h) {
    dense_grad& p = *(new dense_grad(g.in_size, h.out_size));
    p.allocate();
    memset(p.vector, '\0', sizeof(dtype) * p.in_size * p.out_size);
    for(int i = 0; i < h.out_size; i ++)
        for(int j = 0; j < h.in_size; j ++)
            for(int k = 0; k < g.in_size; k ++) {
                p.vector[i * g.in_size + k] += g.get(k, j) * h.get(j, i);
            }
    return &p;
}

void test(grad& g, grad& h) {
    grad& p = *g.get_combine_dst(h);
    std::cout << '[' << name(g) << "]*[" << name(h) << "]=[" << name(p) << ']' << std::endl;
    g.combine(h, p);
    grad& p2 = *general_combine(g, h);
    bool problem = false;
    for(int i = 0; i < p.out_size; i ++) {
        for(int j = 0; j < p.in_size; j ++) {
            if(p2.get(j, i) != p.get(j, i)) problem = true;
        }
    }
    if(problem) std::cout << "problem" << std::endl;
    delete &p;
    delete &p2;
}

int main() {
    dense_grad g1(3, 4);
    g1.vector = new dtype[12] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
    lincomb_grad g2(2, 4);
    g2.vector = new dtype[2] {1.1, 2.1};
    lincomb_modified_grad g3(2, 4);
    g3.vector = new dtype[8] {1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2};
    scalar_grad g4(4);
    g4.vector = new dtype[4] {1.3, 2.3, 3.3, 4.3};

    dense_grad h1(4, 2);
    h1.vector = new dtype[8] {1, 2, 3, 4, 5, 6, 7, 8};
    lincomb_grad h2(2, 2);
    h2.vector = new dtype[2] {1.1, 2.1};
    lincomb_modified_grad h3(2, 2);
    h3.vector = new dtype[4] {1.2, 2.2, 3.2, 4.2};
    scalar_grad h4(4);
    h4.vector = new dtype[4] {1.3, 2.3, 3.3, 4.3};

    test(g1, h1); test(g1, h2); test(g1, h3); test(g1, h4);
    test(g2, h1); test(g2, h2); test(g2, h3); test(g2, h4);
    test(g3, h1); test(g3, h2); test(g3, h3); test(g3, h4);
    test(g4, h1); test(g4, h2); test(g4, h3); test(g4, h4);

    return 0;
}
